#include <iostream>
#include <sstream>
#include <string>

#include <XnCppWrapper.h>

#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>

#include <pcl/io/pcd_io.h>
#include <pcl/point_types.h>
#include <pcl/point_cloud.h>

#include <pcl/kdtree/impl/kdtree_flann.hpp>
#include <pcl/search/impl/organized.hpp>
#include <pcl/filters/impl/statistical_outlier_removal.hpp>
#include <pcl/filters/impl/radius_outlier_removal.hpp>


#define BUFFERSIZE 100

static cv::Mat RGBBuffer[BUFFERSIZE], DepthBuffer[BUFFERSIZE];
static unsigned int buffer_cnt = 0;
static unsigned int frame_cnt = 0;
static unsigned int view_cnt = 1;
stringstream savename;

static XnStatus eResult = XN_STATUS_OK;
static xn::Context mContext;
static xn::DepthGenerator mDepthGenerator;
static xn::ImageGenerator mImageGenerator;

void CheckOpenNIError(XnStatus eResult, std::string sStatus)
{
	if (eResult != XN_STATUS_OK)
		std::cerr << sStatus << " Error: " << xnGetStatusString(eResult) << std::endl;
}

void DeviceInit()
{
	// initial context
	eResult = mContext.InitFromXmlFile("SamplesConfig.xml");
	CheckOpenNIError(eResult, "Initialize context");

	// create depth & image generator
	eResult = mContext.FindExistingNode(XN_NODE_TYPE_DEPTH, mDepthGenerator);
	CheckOpenNIError(eResult, "Create depth generator");

	eResult = mContext.FindExistingNode(XN_NODE_TYPE_IMAGE, mImageGenerator);
	CheckOpenNIError(eResult, "Create image generator");

	// calibration, important!
	mDepthGenerator.GetAlternativeViewPointCap().SetViewPoint(mImageGenerator);
}

void avgImage(cv::Mat& RGB, cv::Mat& Depth)
{
	cv::Mat sumOfRGB(cv::Size(640, 480), CV_16UC3, cv::Scalar::all(0));
	cv::Mat sumOfDepth(cv::Size(640, 480), CV_32S, cv::Scalar::all(0));
	cv::Mat nonNaNOfDepth(cv::Size(640, 480), CV_8U, cv::Scalar::all(0));

	for (size_t i = 0; i < BUFFERSIZE;++i) {
		cv::add(sumOfRGB, RGBBuffer[i], sumOfRGB, cv::noArray(), CV_16UC3);

		for (int h = 0; h < sumOfDepth.rows; ++h) {
			for (int w = 0; w < sumOfDepth.cols; ++w) {
				if (DepthBuffer[i].at<ushort>(h,w) != 0) {
					sumOfDepth.at<unsigned int>(h,w) += DepthBuffer[i].at<ushort>(h,w);
					nonNaNOfDepth.at<uchar>(h,w) += 1;
				}
			}
		}
	}

	// Calculate mean
	divide(sumOfRGB, cv::Scalar::all(BUFFERSIZE), RGB, 1, CV_8UC3);

	for (int h = 0; h < Depth.rows; ++h) {
		for (int w = 0; w < Depth.cols; ++w) {
			if (nonNaNOfDepth.at<uchar>(h,w) != 0) {
				Depth.at<ushort>(h,w) = sumOfDepth.at<unsigned int>(h,w) / nonNaNOfDepth.at<uchar>(h,w);
			}
		}
	}
}

void eliminateNoise(pcl::PointCloud<pcl::PointXYZ>::Ptr &cloud)
{
	pcl::StatisticalOutlierRemoval<pcl::PointXYZ> static_filter;
	static_filter.setInputCloud(cloud);
	static_filter.setMeanK(200);
	static_filter.setStddevMulThresh(1.0);
	static_filter.filter(*cloud);

	pcl::RadiusOutlierRemoval<pcl::PointXYZ> radius_filter;
	radius_filter.setInputCloud(cloud);
	radius_filter.setRadiusSearch(0.01);
	radius_filter.setMinNeighborsInRadius(20);
	radius_filter.filter(*cloud);
}


void saveForeground(cv::Mat& BackgroundRGB, cv::Mat& BackgroundDepth)
{
	float bad_point = std::numeric_limits<float>::quiet_NaN ();
	pcl::PointCloud<pcl::PointXYZRGB>::Ptr cloud (new pcl::PointCloud<pcl::PointXYZRGB>);
	pcl::PointCloud<pcl::PointXYZ>::Ptr saved_points (new pcl::PointCloud<pcl::PointXYZ>);
	cv::Mat RGB_texture(cv::Size(640, 480), CV_8UC3, cv::Scalar::all(0));
	cv::Mat AvgForegroundRGB(cv::Size(640, 480), CV_8UC3, cv::Scalar::all(0));
	cv::Mat AvgForegroundDepth(cv::Size(640, 480), CV_16U, cv::Scalar::all(0));

	avgImage(AvgForegroundRGB, AvgForegroundDepth);

	xn::DepthMetaData depMD;
	mDepthGenerator.GetMetaData(depMD);
	unsigned int uPointNum = depMD.FullXRes() * depMD.FullYRes();

	const XnRGB24Pixel* pImageMap = mImageGenerator.GetRGB24ImageMap();

	cloud->height = depMD.FullYRes();
	cloud->width = depMD.FullXRes();
	cloud->points.resize (cloud->height * cloud->width);

	XnDouble pixel_size;
	XnUInt64 depth_focal_length_SXGA;
	mDepthGenerator.GetRealProperty("ZPPS", pixel_size);
	mDepthGenerator.GetIntProperty ("ZPD", depth_focal_length_SXGA);


	register float constant = 1.0f / ((depth_focal_length_SXGA / pixel_size) * 0.5);
	register int cnt = 0;

	register int centerX = (depMD.FullXRes() >> 1);
	int centerY = (depMD.FullYRes() >> 1);
	for (int v = -centerY; v < centerY; ++v) {
		for (register int u = -centerX; u < centerX; ++u, ++cnt) {
			int raw = v + centerY;
			int col = u + centerX;
			pcl::PointXYZRGB &pt = cloud->points[cnt];

			if (AvgForegroundDepth.at<ushort>(raw, col) == 0) {
				pt.x = pt.y = pt.z = bad_point;
				continue;
			}

			pt.z = AvgForegroundDepth.at<ushort>(raw, col) * 0.001f;
			pt.x = (static_cast<float> (u) * pt.z * 0.1052) / 60;
			pt.y = (static_cast<float> (v) * pt.z * 0.1052) / 60;
			pt.r = pImageMap[cnt].nRed;
			pt.g = pImageMap[cnt].nGreen;
			pt.b = pImageMap[cnt].nBlue;
		}
	}
	//pcl::io::savePCDFileASCII("cloud.pcd", *cloud); // whole point cloud
	//system("pause");

	saved_points->points.resize(cloud->points.size());
	cnt = 0;
	unsigned idx = 0;
	for (unsigned int i = 0; i < depMD.FullYRes(); ++i) {
		for (unsigned int j = 0; j < depMD.FullXRes(); ++j, ++cnt) {
			if (AvgForegroundDepth.at<ushort>(i, j) == 0 || i < 170 || j > 460 || j < 150) // no data or out of range, ignore
				continue;

			if (abs(BackgroundDepth.at<ushort>(i, j) - AvgForegroundDepth.at<ushort>(i, j)) > 15) {
				saved_points->points[idx].x = cloud->points[cnt].x;
				saved_points->points[idx].y = cloud->points[cnt].y;
				saved_points->points[idx].z = cloud->points[cnt].z;
				++idx;

				RGB_texture.at<cv::Vec3b>(i, j) = AvgForegroundRGB.at<cv::Vec3b>(i, j);
			}
		}
	}
	saved_points->width = idx;
	saved_points->height = 1;
	saved_points->resize(saved_points->width * saved_points->height);

	if (saved_points->size() == 0) {
		std::cout << "No foreground saved!!" << std::endl;
	} else {
		eliminateNoise(saved_points);

		savename.str("");
		savename << "RGB_texture_" << view_cnt << ".jpg";
		cv::imwrite(savename.str(), RGB_texture);

		savename.str("");
		savename << "foreground_" << view_cnt << ".pcd";
		pcl::io::savePCDFileASCII(savename.str(), *saved_points);

		std::cout << "Foreground saved\tTotal points: " << saved_points->points.size() << std::endl;
	}
}

int main(int argc, char* argv[])
{
	cv::Mat BackgroundRGB(cv::Size(640, 480), CV_8UC3, cv::Scalar::all(0)), 
			BackgroundDepth(cv::Size(640, 480), CV_16U, cv::Scalar::all(0));

	// buffer init
	for (size_t i = 0; i < BUFFERSIZE; i++) {
		RGBBuffer[i].create(cv::Size(640, 480), CV_8UC3);
		DepthBuffer[i].create(cv::Size(640, 480), CV_16U);
	}

	DeviceInit();

	// read data
	mContext.StartGeneratingAll();
	for(;;) {
		eResult = mContext.WaitAndUpdateAll();
		xn::ImageMetaData imgMD;
		xn::DepthMetaData depMD;
		if (eResult == XN_STATUS_OK) {
			// image
			mImageGenerator.GetMetaData(imgMD);

			cv::Mat tmpImg(imgMD.FullYRes(), imgMD.FullXRes(), CV_8UC3, (void*)imgMD.Data());
			cv::cvtColor(tmpImg, RGBBuffer[buffer_cnt], CV_RGB2BGR);
			cv::imshow("Color Image", RGBBuffer[buffer_cnt]);

			// depth
			mDepthGenerator.GetMetaData(depMD);

			cv::Mat tmpDep(depMD.FullYRes(), depMD.FullXRes(), CV_16U, (void*)depMD.Data());
			tmpDep.copyTo(DepthBuffer[buffer_cnt]);
			//tmpDep.convertTo(DepthBuffer[buffer_cnt], CV_32FC1, 0.0001);
			//cv::imshow("Depth Image", DepthBuffer[buffer_cnt] * 255);
		}

		char c = cv::waitKey(30);
		switch (c) {
			case 27:
				break;
			case 'b':
				avgImage(BackgroundRGB, BackgroundDepth);
				std::cout << "Background saved." << std::endl;
				break;
			case 'f':
				saveForeground(BackgroundRGB, BackgroundDepth);

				savename.str("");
				savename << "RGB_" << view_cnt << ".jpg";
				cv::imwrite(savename.str(), RGBBuffer[buffer_cnt]);

				savename.str("");
				savename << "Depth_" << view_cnt << ".jpg";
				cv::imwrite(savename.str(), DepthBuffer[buffer_cnt] / 3500 * 255);

				++view_cnt;
				break;
			default:
				break;
		}

		++buffer_cnt;
		if (buffer_cnt == BUFFERSIZE)
			buffer_cnt = 0;

		cout << "Frame count: " << frame_cnt++ << "\r";
	}
	// stop
	mContext.StopGeneratingAll();
	mContext.Release();

	return 0;
}